package com.zhhe.tree;

import lombok.Data;

/*
 *Author:ZouHeng
 *Des:红黑树（左倾红黑树）
 *Date:2021-01-18  13:17
 */
public class RedBlackTree<Key extends Comparable<Key>, Value> {

    private Node<Key, Value> root;
    private int size;

    /**
     * 红色标识
     */
    private static final boolean RED = true;

    /**
     * 黑色标识
     */
    private static final boolean BLACK = false;

    public RedBlackTree() {
        this.size = 0;
    }

    /**
     * 判断当前节点的父指向链接是否为红色
     * @return
     */
    private boolean isRed(Node<Key, Value> node) {
        if (node == null) {
            return false;
        }
        return node.isColor();
    }

    /**
     * 左旋调整
     * @param node
     * @return
     */
    private Node<Key, Value> rotateLeft(Node<Key, Value> node) {
        //1、获取该节点的右孩子结点
        Node<Key, Value> rcNode = node.getRight();
        //2、获取该节点的右孩子结点的左孩子结点
        Node<Key, Value> rclcNode = rcNode.getLeft();
        //3、将 该节点的右孩子结点的左孩子结点 设置为该结点的右孩子结点
        node.setRight(rclcNode);
        //4、将该结点设置为该结点右孩子结点的左子结点
        rcNode.setLeft(node);
        //5、将rcNode的颜色变成node的颜色（变为黑色）
        rcNode.setColor(node.isColor());
        //6、将node的颜色变成红色
        node.setColor(RED);
        return rcNode;
    }

    /**
     * 右旋调整
     * @param node
     * @return
     */
    private Node<Key, Value> rotateRight(Node<Key, Value> node) {
        //1、获取该结点的左孩子结点lcNode
        Node<Key, Value> lcNode = node.getLeft();
        //2、获取该结点的左孩子结点的右孩子结点lclcNode
        Node<Key, Value> lcrcNode = lcNode.getRight();
        //3、将该结点的左孩子结点的右孩子结点设置为该结点的左孩子结点（将lcrcNode设为Node的左孩子结点）
        node.setLeft(lcrcNode);
        //4、将该结点设置为该结点左孩子结点的右孩子结点(将node设为lcNode的右孩子结点)
        lcNode.setRight(node);
        //5、将lcNode的颜色变为node的颜色（变为黑色）
        lcNode.setColor(node.isColor());
        //6、将node的颜色设置为红色
        node.setColor(RED);
        return lcNode;
    }

    /**
     * 变色
     * 颜色反转，相当于拆分4-结点
     * @param node
     */
    private void flipColor(Node<Key, Value> node) {
        node.setColor(RED);
        node.getRight().setColor(BLACK);
        node.getLeft().setColor(BLACK);
    }

    /**
     * 获取树中元素的个数
     *
     * @return
     */
    public int size() {
        return this.size;
    }

    public Value get(Key key) {
        return get(root, key);
    }

    /**
     * 从指定的树中，找到key对应的值
     * 该方法与2-3查找树的get方法一样
     * @param node
     * @param key
     * @return
     */
    private Value get(Node<Key, Value> node, Key key) {
        if (node == null) {
            return null;
        }
        int cmp = key.compareTo(node.getKey());
        if (cmp < 0) {
            return get(node.getLeft(), key);
        } else if (cmp > 0) {
            return get(node.getRight(), key);
        } else {
            return node.getValue();
        }
    }

    public void put(Key key, Value value) {
        root = put(root, key, value);
        root.setColor(BLACK);
    }

    public Node<Key, Value> put(Node<Key, Value> node, Key key, Value value) {
        if (node == null) {
            this.size++;
            return new Node<>(null, null, key, value, RED);
        }
        int cmp = key.compareTo(node.getKey());
        if (cmp < 0) {
            node.left = put(node.getLeft(), key, value);
        } else if (cmp > 0) {
            node.right = put(node.getRight(), key, value);
        } else {
            node.setKey(key);
            node.setValue(value);
        }

        //如果当前结点的左链接是黑链接，右链接是红链接，则进行左旋
        if (!isRed(node.getLeft()) && isRed(node.getRight())) {
            node = rotateLeft(node);
        }

        //如果当前结点的左连接是红链接，当前结点左孩子结点的左链接是红链接，则进行右旋
        if (isRed(node.getLeft()) && isRed(node.getLeft().getLeft())) {
            node = rotateRight(node);
        }

        //如果当前结点的左链接是红链接，右链接是红链接，则进行变色
        if (isRed(node.getLeft()) && isRed(node.getRight())) {
            flipColor(node);
        }

        return node;
    }

    @Data
    private static class Node<Key extends Comparable<Key>, Value> {
        private Node<Key, Value> left;
        private Node<Key, Value> right;
        private Key key;
        private Value value;

        /**
         * 如果指向它的链接是红色的，那么该变量的值为true
         * 如果指向它的链接是黑色的，那么该变量的值为false
         */
        private boolean color;

        public Node(Node<Key, Value> left, Node<Key, Value> right, Key key, Value value, boolean color) {
            this.left = left;
            this.right = right;
            this.key = key;
            this.value = value;
            this.color = color;
        }
    }

    public static void main(String[] args) {
        RedBlackTree<Integer, String> tree = new RedBlackTree<>();
        tree.put(4, "二哈");
        tree.put(1, "张三");
        tree.put(3, "李四");
        tree.put(5, "王五");
        System.out.println(tree.size());
    }
}
